#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RocketMQ 异步生产者示例
演示如何使用异步方式发送消息
"""

import time
import json
import logging
from rocketmq.client import Producer, Message, SendStatus
from config import ROCKETMQ_CONFIG, MESSAGE_CONFIG, LOGGING_CONFIG

# 配置日志
logging.basicConfig(
    level=getattr(logging, LOGGING_CONFIG["level"]), 
    format=LOGGING_CONFIG["format"]
)
logger = logging.getLogger(__name__)


class AsyncCallback:
    """异步发送回调类"""
    
    def __init__(self, message_id):
        self.message_id = message_id
        self.start_time = time.time()
    
    def on_success(self, send_result):
        """发送成功回调"""
        elapsed_time = (time.time() - self.start_time) * 1000  # 转换为毫秒
        logger.info(f"✓ 异步消息发送成功 (耗时: {elapsed_time:.2f}ms)")
        logger.info(f"  - 消息ID: {self.message_id}")
        logger.info(f"  - MsgId: {send_result.msg_id}")
        
        # 兼容性处理
        if hasattr(send_result, 'offset'):
            logger.info(f"  - 偏移量: {send_result.offset}")
        if hasattr(send_result, 'status'):
            logger.info(f"  - 状态: {send_result.status}")
    
    def on_exception(self, e):
        """发送异常回调"""
        elapsed_time = (time.time() - self.start_time) * 1000
        logger.error(f"✗ 异步消息发送失败 (耗时: {elapsed_time:.2f}ms)")
        logger.error(f"  - 消息ID: {self.message_id}")
        logger.error(f"  - 错误: {e}")


class AsyncRocketMQProducer:
    """异步 RocketMQ 生产者类"""
    
    def __init__(self):
        self.producer = None
        self.namesrv_addr = ROCKETMQ_CONFIG["namesrv_addr"]
        self.access_key = ROCKETMQ_CONFIG["access_key"]
        self.secret_key = ROCKETMQ_CONFIG["secret_key"]
        self.group_name = ROCKETMQ_CONFIG["producer_group"]
        self.topic = ROCKETMQ_CONFIG["default_topic"]
        self.sent_count = 0
        self.success_count = 0
        self.failed_count = 0
    
    def start(self):
        """启动生产者"""
        try:
            logger.info("=== 启动异步 RocketMQ 生产者 ===")
            logger.info(f"NameServer: {self.namesrv_addr}")
            logger.info(f"生产者组: {self.group_name}")
            logger.info(f"主题: {self.topic}")
            
            # 创建生产者
            self.producer = Producer(self.group_name)
            self.producer.set_name_server_address(self.namesrv_addr)
            self.producer.set_session_credentials(self.access_key, self.secret_key, "ALIYUN")
            
            # 启动生产者
            self.producer.start()
            logger.info("✓ 异步生产者启动成功")
            return True
            
        except Exception as e:
            logger.error(f"✗ 异步生产者启动失败: {e}")
            return False
    
    def send_message_async(self, message_data, tags=None, keys=None, message_id=None):
        """异步发送单条消息"""
        try:
            # 创建消息
            msg = Message(self.topic)
            
            # 设置消息内容
            if isinstance(message_data, dict):
                body = json.dumps(message_data, ensure_ascii=False)
            else:
                body = str(message_data)
            
            msg.set_body(body.encode(MESSAGE_CONFIG["encoding"]))
            
            # 设置消息标签
            if tags:
                msg.set_tags(tags)
            else:
                msg.set_tags(MESSAGE_CONFIG["default_tags"])
            
            # 设置消息键
            if keys:
                msg.set_keys(keys)
            
            # 创建回调对象
            if message_id is None:
                message_id = f"async_msg_{self.sent_count + 1}"
            
            callback = AsyncCallback(message_id)
            
            # 异步发送消息
            self.producer.send_async(msg, callback)
            
            self.sent_count += 1
            logger.info(f"异步消息已提交: {message_id}")
            
            return True
            
        except Exception as e:
            logger.error(f"✗ 异步消息提交失败: {e}")
            self.failed_count += 1
            return False
    
    def send_batch_async(self, message_list, batch_size=5):
        """批量异步发送消息"""
        logger.info(f"=== 批量异步发送 {len(message_list)} 条消息 ===")
        
        total_batches = (len(message_list) + batch_size - 1) // batch_size
        
        for batch_idx in range(total_batches):
            start_idx = batch_idx * batch_size
            end_idx = min(start_idx + batch_size, len(message_list))
            batch_messages = message_list[start_idx:end_idx]
            
            logger.info(f"发送批次 {batch_idx + 1}/{total_batches} ({len(batch_messages)} 条消息)")
            
            for i, msg_data in enumerate(batch_messages):
                message_id = f"batch_{batch_idx + 1}_msg_{i + 1}"
                
                self.send_message_async(
                    message_data=msg_data,
                    tags=f"async_batch_{batch_idx + 1}_{i + 1}",
                    keys=f"async_key_{batch_idx + 1}_{i + 1}",
                    message_id=message_id
                )
            
            # 批次间隔
            if batch_idx < total_batches - 1:
                time.sleep(1)
        
        logger.info(f"批量异步发送完成: 已提交 {self.sent_count} 条消息")
    
    def wait_for_completion(self, timeout=30):
        """等待所有异步消息完成"""
        logger.info(f"等待异步消息完成 (超时: {timeout}秒)...")
        
        start_time = time.time()
        last_count = 0
        
        while time.time() - start_time < timeout:
            time.sleep(1)
            
            # 每5秒打印一次状态
            elapsed = time.time() - start_time
            if int(elapsed) % 5 == 0 and elapsed > 0:
                logger.info(f"异步发送状态: 已提交 {self.sent_count} 条消息")
        
        logger.info("异步发送等待完成")
    
    def get_statistics(self):
        """获取统计信息"""
        return {
            "sent_count": self.sent_count,
            "success_count": self.success_count,
            "failed_count": self.failed_count,
            "pending_count": self.sent_count - self.success_count - self.failed_count
        }
    
    def shutdown(self):
        """关闭生产者"""
        if self.producer:
            # 等待一段时间让异步消息完成
            self.wait_for_completion(10)
            
            self.producer.shutdown()
            logger.info("异步生产者已关闭")


def create_async_demo_messages():
    """创建异步演示消息"""
    messages = []
    
    # 不同类型的消息
    message_types = [
        {"type": "async_text", "content": "Hello Async RocketMQ!"},
        {"type": "async_json", "data": {"user": "async_demo", "action": "send_async"}},
        {"type": "async_number", "value": 100},
        {"type": "async_list", "items": [10, 20, 30, 40, 50]},
        {"type": "async_boolean", "flag": False},
    ]
    
    for i in range(20):  # 发送20条消息
        msg_type = message_types[i % len(message_types)]
        msg_data = {
            "id": i + 1,
            "timestamp": int(time.time()),
            "message": f"异步演示消息 #{i + 1}",
            "data": msg_type,
            "source": "python_async_demo",
            "send_type": "async"
        }
        messages.append(msg_data)
    
    return messages


def main():
    """主函数"""
    logger.info("开始 RocketMQ 异步生产者演示...")
    
    # 创建异步生产者
    producer = AsyncRocketMQProducer()
    
    try:
        # 启动生产者
        if not producer.start():
            return
        
        # 创建演示消息
        messages = create_async_demo_messages()
        
        # 批量异步发送消息
        producer.send_batch_async(messages, batch_size=5)
        
        # 等待异步消息完成
        producer.wait_for_completion(30)
        
        # 获取统计信息
        stats = producer.get_statistics()
        
        # 总结
        logger.info("=== 异步演示总结 ===")
        logger.info(f"提交消息数: {stats['sent_count']}")
        logger.info(f"成功消息数: {stats['success_count']}")
        logger.info(f"失败消息数: {stats['failed_count']}")
        logger.info(f"待处理消息数: {stats['pending_count']}")
        
        if stats['sent_count'] > 0:
            logger.info("✓ 异步生产者演示成功完成")
        else:
            logger.error("✗ 异步生产者演示失败")
            
    except KeyboardInterrupt:
        logger.info("用户中断演示")
    except Exception as e:
        logger.error(f"演示过程中发生错误: {e}")
    finally:
        producer.shutdown()


if __name__ == "__main__":
    main() 